#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/4/9 15:29
# @Author  : zzp
# @File    : RemoteWorker
# @Software: PyCharm
import re
import subprocess

from PyQt5.QtCore import pyqtSignal, QThread


class TrainWorker(QThread):
    output_signal = pyqtSignal(str)
    progress_signal = pyqtSignal(int)

    def __init__(self, params):
        super().__init__()
        self.proc = None
        self.cmd = None
        self.conda_env = "python39"
        self.params = params
        self._running = True

    def build_cmd(self):
        # 1. 构造命令行参数字符串
        script = self.params.get("script", "train.py")
        args_str = " ".join(
            f"--{key} {value}" for key, value in self.params.items() if key != "script"
        )

        # 2. 构造完整命令（使用 conda run 激活虚拟环境）
        cmd = (
            f"conda activate {self.conda_env} && python {script} {args_str}"
        )
        self.output_signal.emit(f"[执行命令] {cmd}")
        return cmd

    def run(self):
        try:
            self.cmd = self.build_cmd()
            self.proc = subprocess.Popen(
                self.cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                shell=True,
                text=True,
            )
            for line in iter(self.proc.stdout.readline, ""):
                if line:
                    self.output_signal.emit(line.strip())
                # 进度条更新代码 —— 开始
                # 在这里我们提取 'Test: [20/80]' 格式的进度信息
                match = re.search(r"Test: \[(\d+)/(\d+)\]", line)
                if match:
                    completed = int(match.group(1))
                    total = int(match.group(2))
                    if total > 0:
                        percent = int((completed / total) * 100)
                        self.progress_signal.emit(percent)
                # 进度条更新代码 —— 结束
            if self.proc.returncode != 0:
                self.output_signal.emit(f"[错误] 进程退出，返回码：{self.proc.returncode}")
            else:
                self.output_signal.emit("[完成] 命令执行完毕。")
        except Exception as e:
            self.output_signal.emit(f"[异常] 执行命令时发生错误：{str(e)}")
        finally:
            self.release_resources()

    def release_resources(self):
        # 进程关闭
        self.output_signal.emit("进程关闭中")
        if self.proc:
            self.proc.terminate()  # 发送SIGTERM信号
            # self.proc.wait()  # 等待进程终止
        self.output_signal.emit("进程关闭完成")
        self._running = False
